package com.zhuangjie.mvc.ws.core;

import com.zhaungjie.base.utils.LambdaStructureParser.LambdaStructureParser;
import com.zhaungjie.base.utils.ReflectionUtils;
import com.zhaungjie.base.utils.TypeConvertCompatibleUtils;
import com.zhuangjie.mvc.ws.Interceptor.WebSocketInterceptor;
import com.zhuangjie.mvc.ws.annotation.*;
import com.zhuangjie.mvc.ws.constant.WSConstant;
import com.zhuangjie.mvc.ws.entity.WSEventData;
import com.zhuangjie.mvc.ws.entity.WSMessageSessionInfo;
import com.zhuangjie.mvc.ws.entity.WSResult;
import com.zhuangjie.mvc.ws.utils.ContextPathHelper;
import com.zhuangjie.mvc.ws.utils.PathVarUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.lang.annotation.Annotation;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 抽象WS service
 *
 * @author zhuangjie
 * @date 2024/01/12
 */

@Component
@Aspect
public class WebSocketDispatch extends TextWebSocketHandler {

    private Map<String, Object> listenerMap = new ConcurrentHashMap<>();

    private Map<String, Object> putMessageMap = new ConcurrentHashMap<>();

    @Autowired
    private WSEvent wsEvent;

    @Autowired
    private ApplicationContext applicationContext;

    @Autowired
    private WebSocketSessionStore sessionStore;

    @Autowired
    private ContextPathHelper contextPathHelper;

    @Autowired
    private WebSocketInterceptorStore interceptorStore;

    /**
     * 列出注册数据
     *
     * @return {@link Set}<{@link String}>
     */
    public Set<String> initAndGetDispatchHandlerRegisterInfo() {
        // 获取标有YourCustomAnnotation注解的所有beans
        Set<String> registerArray = new HashSet<>();
        Map<String, Object> listenerBean = applicationContext.getBeansWithAnnotation(WSListener.class);
        Map<String, Object> wsPutBean = applicationContext.getBeansWithAnnotation(WSPut.class);
        if (listenerBean != null) {
            this.listenerMap = listenerBean.values().stream().collect(Collectors.toMap(bean -> bean.getClass().getAnnotation(WSListener.class).ws(), bean -> bean));
        }
        if (wsPutBean != null) {
            this.putMessageMap = wsPutBean.values().stream().collect(Collectors.toMap(bean -> bean.getClass().getAnnotation(WSPut.class).ws(), bean -> bean));
        }
        registerArray.addAll(this.listenerMap.keySet());
        registerArray.addAll(this.putMessageMap.keySet());
        return registerArray;
    }


    /**
     * 拦截标记了指定注解的方法的调用（消息发送）
     */
    @Around("@annotation(com.zhuangjie.mvc.ws.annotation.MessagePut)")
    public Object onMessagePut(ProceedingJoinPoint pjp) throws Throwable {
        // 获取AOP拦截的方法
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();

        String rootPath = findPutRootPath(method);
        String requestPath = findPutRequestPath(method);
        // 执行方法，获取返回的ws数据
        Object message = pjp.proceed();
        // 获取要发送的给谁（serviceIds）
        String[] sessionIds = findRecipientServiceSessionIds(method, pjp);
        // 消息组装
        WSResult wsResult = new WSResult();
        wsResult.setMessageType(requestPath);
        wsResult.setData(message);
        // 当前节点-发送
        List<String> unprocessableIds = sessionStore.sendMessage(contextPathHelper.fullRequestRootPath(rootPath), wsResult, sessionIds);
        // 当存在无法处理的时，通知给其它节点来处理（分布式WS核心）
        if (unprocessableIds != null && !unprocessableIds.isEmpty()) {
            WSEventData wsEventData = new WSEventData();
            wsEventData.setKey(contextPathHelper.fullRequestRootPath(rootPath));
            wsEventData.setIds(unprocessableIds.toArray(new String[unprocessableIds.size()]));
            wsEventData.setData(wsResult);
            wsEvent.emit(wsEventData);
        }
        return message;
    }

    /**
     * 查找收件人服务会话ID
     *
     * @param method 方法
     * @param pjp    pjp
     * @return {@link String[]}
     */
    private String[] findRecipientServiceSessionIds(Method method, ProceedingJoinPoint pjp) {
        // 看调用的参数中，哪个参数上有@Send注解
        int index = ReflectionUtils.getOneParamIndexByAnnotation(method, Send.class);
        String[] sessionIds;
        if (index >= 0) {
            Class<?> sendAnnotationModificationParameterType = method.getParameterTypes()[index];
            if (!(sendAnnotationModificationParameterType.equals(String.class) || sendAnnotationModificationParameterType.equals(String[].class))) {
                throw new RuntimeException("@Send标注的属性是一个String类型或String数组！");
            }
            Object sendAnnotationModificationParameterValue = pjp.getArgs()[index];
            if (sendAnnotationModificationParameterValue == null) {
                sessionIds = new String[]{};
            } else if (sendAnnotationModificationParameterValue instanceof String) {
                sessionIds = new String[]{(String) sendAnnotationModificationParameterValue};
            } else {
                sessionIds = (String[]) sendAnnotationModificationParameterValue;
            }
        } else {
            // 当没有标注@Send的方法参数时，通知对应session集中的所有会话
            sessionIds = new String[]{WSConstant.ANY_SESSION_ID};
        }
        return sessionIds;
    }

    private String findPutRequestPath(Method method) throws Exception {
        MessagePut messagePut = method.getAnnotation(MessagePut.class);
        if (messagePut == null) {
            // 当没有标注@WSPut时，获取失败，抛出异常
            throw new Exception("没有标注@MessagePut，获取失败");
        }
        return messagePut.value();
    }

    private String findPutRootPath(Method method) throws Exception {
        Class<?> declaringClass = method.getDeclaringClass();
        WSPut wsPut = declaringClass.getAnnotation(WSPut.class);
        if (wsPut == null) {
            // 当没有标注@WSPut时，获取失败，抛出异常
            throw new Exception("没有标注@WSPut，获取失败");
        }
        return wsPut.ws();
    }


    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        WebSocketInterceptor interceptor = interceptorStore.getInterceptor(session);
        if (interceptor == null) {
            throw new Exception("无法确认Session身份，请配置拦截器拦截设置。");
        }
        String serviceId = interceptor.onRegister(session);

        sessionStore.put(serviceId, session);
        // onOpen
        invokeWSHandler(session,(object)->{
            List<Method> methods = ReflectionUtils.findAnnotatedMethodsInClass(object.getClass(), OnOpen.class);
            if (methods == null || methods.isEmpty()) {
                return null;
            }
            return methods.get(0);
        },(handler)->genWSMessageHandlerMethodParamValues(handler,null,session,new WSMessageSessionInfo(serviceId,session)));
    }

    /**
     * 未找到上处理程序
     */
    public void onHandlerNotFound(WebSocketSession session, TextMessage message) {
        throw new RuntimeException("没有匹配的Handler，消息无法调度处理！");
    }

    /**
     * 处理文本消息
     *
     * @param session     一场
     * @param textMessage 短信
     * @throws Exception 例外
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage textMessage) throws Exception {
        invokeWSHandler(session, (foundObject) -> {
            TextMessageParser.TextMessageBody<Object> messageBody = TextMessageParser.parse(textMessage, Object.class);
            Method handler = null;
            if (foundObject != null) handler = findHandler(foundObject, messageBody.getRequest());
            if (handler == null) {
                List<Method> methods = ReflectionUtils.findAnnotatedMethodsInClass(foundObject.getClass(), OnHandlerNotFound.class);
                handler = methods == null || methods.isEmpty() ? LambdaStructureParser.convertToMethod(this::onHandlerNotFound) : methods.get(0);
            }
            return handler;
        }, (handler) -> {
            WSMessageSessionInfo sessionInfo = new WSMessageSessionInfo(sessionStore.findServiceId(session), session);
            return genWSMessageHandlerMethodParamValues(handler,textMessage,session, textMessage, sessionInfo);
        });
    }

    /**
     * 调用Web Socket处理程序Handler
     *
     * @param session             一场
     * @param methodSelector      方法选择器
     * @param paramValuesSupplier 供应商参数值
     * @return {@link Object}
     * @throws InvocationTargetException 调用目标异常
     * @throws IllegalAccessException    非法访问异常
     */
    private Object invokeWSHandler(WebSocketSession session, Function<Object, Method> methodSelector, Function<Method, Object[]> paramValuesSupplier) throws InvocationTargetException, IllegalAccessException {
        String requestRootPath = session.getUri().getPath();
        Object objectOnHandler = listenerMap.get(contextPathHelper.removeContextPath(requestRootPath));
        Method handler = methodSelector.apply(objectOnHandler);
        if (handler == null) throw new RuntimeException("WebSocket处理程序未找到");
        Object[] parameters = paramValuesSupplier.apply(handler);
        // 反射调用
        handler.setAccessible(true);
        return handler.invoke(objectOnHandler, parameters);
    }


    /**
     * 生成WebSocket的消息处理器Handler方法参数值
     *
     * @param handlerMethod     handler方法
     * @param byClassFillValues 填充值
     * @return {@link Object[]}
     */
    public Object[] genWSMessageHandlerMethodParamValues(Method handlerMethod,TextMessage textMessage,  Object... byClassFillValues) {
        Map<Class, ReflectionUtils.ParameterValueProviderByClass> byClassRuleMap = new HashMap<>();
        Map<Class<? extends Annotation>, ReflectionUtils.ParameterValueProviderByAnnotation> byAnnotaationRuleMap = new HashMap<>();
        for (Object fillValue : byClassFillValues) {
            byClassRuleMap.put(fillValue.getClass(), (annotations, clazz) -> fillValue);
        }
        byAnnotaationRuleMap.put(MessageBody.class, (annotation, clazz) -> {
            return TextMessageParser.parse(textMessage, clazz).getMessage();
        });
        byAnnotaationRuleMap.put(MessagePathVariable.class, (annotation, clazz) -> {
            String pathVarName = ((MessagePathVariable) annotation).value();
            OnMessage onMessage = handlerMethod.getAnnotation(OnMessage.class);
            String handlerRequestPath = onMessage.value();
            TextMessageParser.TextMessageBody<?> messageBody = TextMessageParser.parse(textMessage, clazz);
            Map<String, String> pathVars = PathVarUtils.getPathVars(handlerRequestPath, messageBody.getRequest());
            return TypeConvertCompatibleUtils.convert(pathVars.get(pathVarName), clazz);
        });
        return ReflectionUtils.genMethodParamValues(handlerMethod, byClassRuleMap, byAnnotaationRuleMap);
    }


    /**
     * 查找来消息-处理程序
     *
     * @param objectOnHandler      处理程序上对象
     * @param messageRequestSuffix 消息请求后缀
     * @return {@link Method}
     */
    private Method findHandler(Object objectOnHandler, String messageRequestSuffix) {
        // 从所在类中获取目标method
        List<Method> methods = ReflectionUtils.getAllMethods(objectOnHandler.getClass());
        for (Method method : methods) {
            OnMessage onMessageAnnotation = method.getAnnotation(OnMessage.class);
            if (onMessageAnnotation == null) continue;
            String requestPathOnMethod = onMessageAnnotation.value();
            if (Objects.equals(requestPathOnMethod, messageRequestSuffix) || PathVarUtils.isMate(requestPathOnMethod, messageRequestSuffix)) {
                return method;
            }
        }
        return null;
    }

//    private Object[] fillParameters(Method handler, WebSocketSession session, TextMessage textMessage) throws InvocationTargetException, InstantiationException, IllegalAccessException, NoSuchMethodException {
//        // 获取所有参数类型
//        Class<?>[] parameterTypes = handler.getParameterTypes();
//        Annotation[][] parameterAnnotations = handler.getParameterAnnotations();
//        Object[] parametersValue = new Object[parameterTypes.length];
//        // 循环所有参数，通过反射创建实例（这里假设都有默认构造函数）
//        for (int i = 0; i < parameterTypes.length; i++) {
//            Class<?> parameterType = parameterTypes[i];
//            if (parametersValue != null && parameterType.isInstance(session)) {
//                parametersValue[i] = session;
//            } else if (parameterType != null && parameterType.isInstance(textMessage)) {
//                parametersValue[i] = textMessage;
//            } else if (isHasParamsAnnotation(parameterAnnotations[i], MessageBody.class)) {
//                TextMessageParser.TextMessageBody<?> messageBody = TextMessageParser.parse(textMessage, parameterType);
//                parametersValue[i] = messageBody.getMessage();
//            } else if(isHasParamsAnnotation(parameterAnnotations[i], MessagePathVariable.class)) {
//                TextMessageParser.TextMessageBody<?> messageBody = TextMessageParser.parse(textMessage, parameterType);
//                OnMessage onMessage = handler.getAnnotation(OnMessage.class);
//                String handlerRequestPath = onMessage.value();
//                Map<String, String> pathVars = PathVarUtils.getPathVars(handlerRequestPath, messageBody.getRequest());
//                MessagePathVariable messagePathVariable = getParamsAnnotation(parameterAnnotations[i], MessagePathVariable.class);
//                parametersValue[i] = TypeConvertCompatibleUtils.convert(pathVars.get(messagePathVariable.value()),parameterType);
//            } else if (WSMessageSessionInfo.class.equals(parameterType)){
//                WSMessageSessionInfo sessionInfo = new WSMessageSessionInfo();
//                sessionInfo.setServiceId(sessionStore.findServiceId(session));
//                sessionInfo.setSession(session);
//                sessionInfo.setMessage(textMessage);
//                parametersValue[i] = sessionInfo;
//            }else {
//                parametersValue[i] = null;
//            }
//        }
//        return parametersValue;
//    }

//    private<T extends Annotation> boolean isHasParamsAnnotation(Annotation[] parameterAnnotation, Class<T> clazz) {
//        for (Annotation annotation : parameterAnnotation) {
//            if (clazz.isInstance(annotation)) {
//                return true;
//            }
//        }
//        return false;
//    }
//    private<T extends Annotation> T getParamsAnnotation(Annotation[] parameterAnnotation, Class<T> clazz) {
//        for (Annotation annotation : parameterAnnotation) {
//            if (clazz.isInstance(annotation)) {
//                return (T)annotation;
//            }
//        }
//        return null;
//    }


    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        String serviceId = sessionStore.findServiceId(session);
        sessionStore.remove(session);
        invokeWSHandler(session,(object)->{
            List<Method> methods = ReflectionUtils.findAnnotatedMethodsInClass(object.getClass(), OnError.class);
            if (methods == null || methods.isEmpty()) {
                return null;
            }
            return methods.get(0);
        },(handler)->genWSMessageHandlerMethodParamValues(handler,null,session,exception,new WSMessageSessionInfo(serviceId,session)));
    }


    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String serviceId = sessionStore.findServiceId(session);
        sessionStore.remove(session);
        invokeWSHandler(session,(object)->{
            List<Method> methods = ReflectionUtils.findAnnotatedMethodsInClass(object.getClass(), OnClose.class);
            if (methods == null || methods.isEmpty()) {
                return null;
            }
            return methods.get(0);
        },(handler)->genWSMessageHandlerMethodParamValues(handler,null,session,status,new WSMessageSessionInfo(serviceId,session)));
    }


}
